#!/usr/bin/env python3
"""
质量感知提示生成器

基于图像质量评估的提示生成模块，优先在高质量区域生成提示
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Dict, List, Tuple, Optional, Any

import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parent.parent.parent))
from .base_prompt import BasePromptGenerator
from utils.logger import get_logger


class QualityAwarePromptGenerator(BasePromptGenerator):
    """质量感知提示生成器"""
    
    def __init__(self, 
                 input_dim: int = 256,
                 prompt_dim: int = 256,
                 quality_threshold: float = 0.5,
                 device: str = "cuda"):
        """
        初始化质量感知提示生成器
        
        Args:
            input_dim: 输入特征维度
            prompt_dim: 提示特征维度
            quality_threshold: 质量阈值
            device: 设备
        """
        super().__init__(input_dim, prompt_dim, device)
        
        self.quality_threshold = quality_threshold
        
        # 质量评估器
        self.quality_assessor = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 128, 3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((16, 16)),
            nn.Flatten(),
            nn.Linear(128 * 16 * 16, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 1),
            nn.Sigmoid()
        )
        
        # 质量特征提取器
        self.quality_feature_extractor = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 128, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(128, 256, 3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((32, 32))
        )
        
        # 质量引导的注意力机制
        self.quality_attention = nn.MultiheadAttention(
            embed_dim=prompt_dim,
            num_heads=8,
            dropout=0.1,
            batch_first=True
        )
        
        # 质量感知的提示生成器
        self.quality_guided_generator = nn.ModuleDict({
            'points': nn.Sequential(
                nn.Linear(prompt_dim + 1, prompt_dim // 2),  # +1 for quality score
                nn.ReLU(),
                nn.Linear(prompt_dim // 2, 2)
            ),
            'boxes': nn.Sequential(
                nn.Linear(prompt_dim + 1, prompt_dim // 2),
                nn.ReLU(),
                nn.Linear(prompt_dim // 2, 4)
            ),
            'masks': nn.Sequential(
                nn.Linear(prompt_dim + 1, prompt_dim // 2),
                nn.ReLU(),
                nn.Linear(prompt_dim // 2, 32 * 32)
            )
        })
        
        # 质量权重网络
        self.quality_weight_net = nn.Sequential(
            nn.Linear(1, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()
        )
        
        self.logger.info(f"质量感知提示生成器初始化完成: {self.get_model_info()}")
    
    def assess_image_quality(self, image: torch.Tensor) -> torch.Tensor:
        """
        评估图像质量
        
        Args:
            image: 输入图像 (B, C, H, W)
            
        Returns:
            质量分数 (B, 1)
        """
        return self.quality_assessor(image)
    
    def extract_quality_features(self, image: torch.Tensor) -> torch.Tensor:
        """
        提取质量相关特征
        
        Args:
            image: 输入图像 (B, C, H, W)
            
        Returns:
            质量特征 (B, C, H, W)
        """
        return self.quality_feature_extractor(image)
    
    def generate_quality_aware_prompts(self, 
                                     features: torch.Tensor,
                                     quality_scores: torch.Tensor) -> Dict[str, torch.Tensor]:
        """
        生成质量感知提示
        
        Args:
            features: 图像特征 (B, C, H, W)
            quality_scores: 质量分数 (B, 1)
            
        Returns:
            提示字典
        """
        batch_size, channels, height, width = features.shape
        
        # 将特征展平为序列
        features_flat = features.view(batch_size, channels, -1).transpose(1, 2)  # (B, H*W, C)
        
        # 扩展质量分数到每个空间位置
        quality_expanded = quality_scores.unsqueeze(1).repeat(1, height * width, 1)  # (B, H*W, 1)
        
        # 融合特征和质量信息
        fused_features = torch.cat([features_flat, quality_expanded], dim=-1)  # (B, H*W, C+1)
        
        # 生成提示
        prompts = {}
        for prompt_type, generator in self.quality_guided_generator.items():
            prompt_logits = generator(fused_features)  # (B, H*W, output_dim)
            
            if prompt_type == 'points':
                # 点提示：选择质量最高的位置
                quality_weights = quality_expanded.squeeze(-1)  # (B, H*W)
                weighted_logits = prompt_logits * quality_weights.unsqueeze(-1)
                prompts[prompt_type] = weighted_logits
                
            elif prompt_type == 'boxes':
                # 框提示：基于质量分数调整框的大小
                quality_weights = self.quality_weight_net(quality_expanded)  # (B, H*W, 1)
                weighted_logits = prompt_logits * quality_weights
                prompts[prompt_type] = weighted_logits
                
            elif prompt_type == 'masks':
                # 掩码提示：质量高的区域权重更大
                quality_weights = quality_expanded.squeeze(-1)  # (B, H*W)
                mask_logits = prompt_logits.view(batch_size, height, width, -1)
                mask_logits = mask_logits * quality_weights.view(batch_size, height, width, 1)
                prompts[prompt_type] = mask_logits.view(batch_size, height * width, -1)
        
        # 添加质量信息
        prompts['quality_scores'] = quality_scores
        prompts['quality_weights'] = quality_expanded
        
        return prompts
    
    def select_high_quality_regions(self, 
                                  features: torch.Tensor,
                                  quality_scores: torch.Tensor,
                                  num_regions: int = 5) -> torch.Tensor:
        """
        选择高质量区域
        
        Args:
            features: 图像特征 (B, C, H, W)
            quality_scores: 质量分数 (B, 1)
            num_regions: 选择的区域数量
            
        Returns:
            高质量区域特征 (B, num_regions, C)
        """
        batch_size, channels, height, width = features.shape
        
        # 计算每个位置的质量分数
        quality_map = self.assess_image_quality(features)  # (B, 1, H, W)
        quality_map = F.interpolate(quality_map, size=(height, width), mode='bilinear', align_corners=False)
        
        # 选择质量最高的区域
        quality_flat = quality_map.view(batch_size, -1)  # (B, H*W)
        _, top_indices = torch.topk(quality_flat, num_regions, dim=1)  # (B, num_regions)
        
        # 提取对应位置的特征
        features_flat = features.view(batch_size, channels, -1)  # (B, C, H*W)
        top_features = torch.gather(features_flat, 2, top_indices.unsqueeze(1).expand(-1, channels, -1))
        top_features = top_features.transpose(1, 2)  # (B, num_regions, C)
        
        return top_features
    
    def forward(self, 
                image_features: torch.Tensor,
                context: Optional[Dict[str, Any]] = None) -> Dict[str, torch.Tensor]:
        """
        前向传播
        
        Args:
            image_features: 图像特征 (B, C, H, W)
            context: 上下文信息
            
        Returns:
            提示字典
        """
        # 评估图像质量
        quality_scores = self.assess_image_quality(image_features)
        
        # 提取质量相关特征
        quality_features = self.extract_quality_features(image_features)
        
        # 生成质量感知提示
        prompts = self.generate_quality_aware_prompts(quality_features, quality_scores)
        
        # 选择高质量区域
        high_quality_regions = self.select_high_quality_regions(quality_features, quality_scores)
        prompts['high_quality_regions'] = high_quality_regions
        
        # 质量引导的注意力
        attended_features, attention_weights = self.quality_attention(
            quality_features.view(quality_features.size(0), -1, quality_features.size(1)),
            quality_features.view(quality_features.size(0), -1, quality_features.size(1)),
            quality_features.view(quality_features.size(0), -1, quality_features.size(1))
        )
        
        prompts['attended_features'] = attended_features
        prompts['attention_weights'] = attention_weights
        
        return prompts
    
    def generate_prompts(self, 
                        image: torch.Tensor,
                        metadata: Optional[Dict[str, Any]] = None) -> Dict[str, torch.Tensor]:
        """
        生成提示
        
        Args:
            image: 输入图像 (B, C, H, W)
            metadata: 元数据
            
        Returns:
            提示字典
        """
        # 提取图像特征
        image_features = self.extract_features(image)
        
        # 生成质量感知提示
        prompts = self.forward(image_features, metadata)
        
        return prompts
    
    def compute_quality_loss(self, 
                           predictions: Dict[str, torch.Tensor],
                           targets: Dict[str, torch.Tensor]) -> torch.Tensor:
        """
        计算质量感知损失
        
        Args:
            predictions: 预测结果
            targets: 目标结果
            
        Returns:
            损失值
        """
        total_loss = 0.0
        
        # 基础提示损失
        total_loss += self.compute_loss(predictions, targets)
        
        # 质量一致性损失
        if 'quality_scores' in predictions and 'quality_scores' in targets:
            quality_loss = F.mse_loss(
                predictions['quality_scores'],
                targets['quality_scores']
            )
            total_loss += 0.1 * quality_loss
        
        # 质量引导损失：高质量区域应该有更好的预测
        if 'quality_weights' in predictions and 'points' in predictions:
            quality_weights = predictions['quality_weights'].squeeze(-1)  # (B, H*W)
            point_predictions = predictions['points']  # (B, H*W, 2)
            
            # 计算质量加权损失
            weighted_loss = torch.mean(
                quality_weights * torch.norm(point_predictions, dim=-1)
            )
            total_loss += 0.01 * weighted_loss
        
        return total_loss
    
    def get_quality_analysis(self, image: torch.Tensor) -> Dict[str, Any]:
        """
        获取质量分析结果
        
        Args:
            image: 输入图像 (B, C, H, W)
            
        Returns:
            质量分析结果
        """
        with torch.no_grad():
            # 评估质量
            quality_scores = self.assess_image_quality(image)
            
            # 提取质量特征
            quality_features = self.extract_quality_features(image)
            
            # 生成质量图
            quality_map = F.interpolate(
                quality_scores, 
                size=(image.size(2), image.size(3)), 
                mode='bilinear', 
                align_corners=False
            )
            
            return {
                'quality_scores': quality_scores.cpu().numpy(),
                'quality_map': quality_map.cpu().numpy(),
                'quality_features': quality_features.cpu().numpy()
            }
